#pragma once
#include <queue>
#include <arm_neon.h> // 引入NEON头文件

std::priority_queue<std::pair<float, uint32_t> > simd_flat(float* base, float* query, size_t base_number, size_t vecdim, size_t k) {
    std::priority_queue<std::pair<float, uint32_t> > q;

    for(int i = 0; i < base_number; ++i) {
        float dis = 0;
        const float* vec = base + i * vecdim; // 当前基准向量指针
        
        // NEON向量化计算部分
        int d = 0;
        float32x4_t sum = vdupq_n_f32(0); // 初始化累加器为0
        
        // 每次处理4个元素（128位寄存器）
        for(; d + 4 <= vecdim; d += 4) {
            // 加载4个float元素
            float32x4_t v_base = vld1q_f32(vec + d);
            float32x4_t v_query = vld1q_f32(query + d);
            
            // 乘加运算：sum += base * query
            sum = vmlaq_f32(sum, v_base, v_query);
        }
        
        // 将NEON寄存器中的4个float结果水平相加
        dis = vaddvq_f32(sum);
        
        // 处理剩余不足4个的元素
        for(; d < vecdim; ++d) {
            dis += vec[d] * query[d];
        }
        
        dis = 1 - dis;

        if(q.size() < k) {
            q.push({dis, i});
        } else {
            if(dis < q.top().first) {
                q.push({dis, i});
                q.pop();
            }
        }
    }
    return q;
}